{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TFX Components Walk-through"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "1.  Develop a high level understanding of TFX pipeline components.\n",
    "2.  Learn how to use a TFX Interactive Context for prototype development of TFX pipelines.\n",
    "3.  Work with the Tensorflow Data Validation (TFDV) library to check and analyze input data.\n",
    "4.  Utilize the Tensorflow Transform (TFT) library for scalable data preprocessing and feature transformations.\n",
    "5.  Employ the Tensorflow Model Analysis (TFMA) library for model evaluation.\n",
    "\n",
    "In this lab, you will work with the [Covertype Data Set](https://github.com/jarokaz/mlops-labs/blob/master/datasets/covertype/README.md) and use TFX to analyze, understand, and pre-process the dataset and train, analyze, validate, and deploy a multi-class classification model to predict the type of forest cover from cartographic features.\n",
    "\n",
    "You will utilize  **TFX Interactive Context** to work with the TFX components interactivelly in a Jupyter notebook environment. Working in an interactive notebook is useful when doing initial data exploration, experimenting with models, and designing ML pipelines. You should be aware that there are differences in the way interactive notebooks are orchestrated, and how they access metadata artifacts. In a production deployment of TFX on GCP, you will use an orchestrator such as Kubeflow Pipelines, or Cloud Composer. In an interactive mode, the notebook itself is the orchestrator, running each TFX component as you execute the notebook cells. In a production deployment, ML Metadata will be managed in a scalabe database like MySQL, and artifacts in apersistent store such as Google Cloud Storage. In an interactive mode, both properties and payloads are stored in a local file system of the Jupyter host.\n",
    "\n",
    "**Setup Note:**\n",
    "Currently, TFMA visualizations do not render properly in JupyterLab. It is recommended to run this notebook in Jupyter Classic Notebook. To switch to Classic Notebook select *Launch Classic Notebook* from the *Help* menu."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import absl\n",
    "import os\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_data_validation as tfdv\n",
    "import tensorflow_model_analysis as tfma\n",
    "import tensorflow_transform as tft\n",
    "import tfx\n",
    "\n",
    "from pprint import pprint\n",
    "from tensorflow_metadata.proto.v0 import schema_pb2, statistics_pb2, anomalies_pb2\n",
    "from tensorflow_transform.tf_metadata import schema_utils\n",
    "from tfx.components import CsvExampleGen\n",
    "from tfx.components import Evaluator\n",
    "from tfx.components import ExampleValidator\n",
    "from tfx.components import InfraValidator\n",
    "from tfx.components import Pusher\n",
    "from tfx.components import ResolverNode\n",
    "from tfx.components import SchemaGen\n",
    "from tfx.components import StatisticsGen\n",
    "from tfx.components import Trainer\n",
    "from tfx.components import Transform\n",
    "from tfx.components import Tuner\n",
    "from tfx.dsl.components.base import executor_spec\n",
    "from tfx.components.common_nodes.importer_node import ImporterNode\n",
    "from tfx.components.trainer import executor as trainer_executor\n",
    "from tfx.dsl.experimental import latest_blessed_model_resolver\n",
    "from tfx.orchestration import metadata\n",
    "from tfx.orchestration import pipeline\n",
    "from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext\n",
    "from tfx.proto import evaluator_pb2\n",
    "from tfx.proto import example_gen_pb2\n",
    "from tfx.proto import infra_validator_pb2\n",
    "from tfx.proto import pusher_pb2\n",
    "from tfx.proto import trainer_pb2\n",
    "from tfx.proto.evaluator_pb2 import SingleSlicingSpec\n",
    "\n",
    "from tfx.types import Channel\n",
    "from tfx.types.standard_artifacts import Model\n",
    "from tfx.types.standard_artifacts import HyperParameters\n",
    "from tfx.types.standard_artifacts import ModelBlessing\n",
    "from tfx.types.standard_artifacts import InfraBlessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: this lab was developed and tested with the following TF ecosystem package versions:\n",
    "\n",
    "`Tensorflow Version: 2.3.1`  \n",
    "`TFX Version: 0.25.0`  \n",
    "`TFDV Version: 0.25.0`  \n",
    "`TFMA Version: 0.25.0`\n",
    "\n",
    "If you encounter errors with the above imports (e.g. TFX component not found), check your package versions in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Tensorflow Version:\", tf.__version__)\n",
    "print(\"TFX Version:\", tfx.__version__)\n",
    "print(\"TFDV Version:\", tfdv.__version__)\n",
    "print(\"TFMA Version:\", tfma.VERSION_STRING)\n",
    "\n",
    "absl.logging.set_verbosity(absl.logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the versions above do not match, update your packages in the current Jupyter kernel below. The default `%pip` package installation location is not on your system installation PATH; use the command below to append the local installation path to pick up the latest package versions. Note that you may also need to restart your notebook kernel to pick up the specified package versions and re-run the imports cell above before proceeding with the lab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PATH'] += os.pathsep + '/home/jupyter/.local/bin'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Optional) Run the cell below only if your package versions do not match the lab defaults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --user tensorflow==2.3.1\n",
    "%pip install --upgrade --user tfx==0.25.0\n",
    "%pip install --upgrade --user tensorflow_data_validation==0.25.0\n",
    "%pip install --upgrade --user tensorflow_model_analysis==0.25.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure lab settings\n",
    "\n",
    "Set constants, location paths and other environment settings. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ARTIFACT_STORE = os.path.join(os.sep, 'home', 'jupyter', 'artifact-store')\n",
    "SERVING_MODEL_DIR=os.path.join(os.sep, 'home', 'jupyter', 'serving_model')\n",
    "DATA_ROOT = 'gs://workshop-datasets/covertype/small'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Interactive Context\n",
    "\n",
    "TFX Interactive Context allows you to create and run TFX Components in an interactive mode. It is designed to support experimentation and development in a Jupyter Notebook environment. It is an experimental feature and major changes to interface and functionality are expected. When creating the interactive context you can specifiy the following parameters:\n",
    "- `pipeline_name` - Optional name of the pipeline for ML Metadata tracking purposes. If not specified, a name will be generated for you.\n",
    "- `pipeline_root` - Optional path to the root of the pipeline's outputs. If not specified, an ephemeral temporary directory will be created and used.\n",
    "- `metadata_connection_config` - Optional `metadata_store_pb2.ConnectionConfig` instance used to configure connection to a ML Metadata connection. If not specified, an ephemeral SQLite MLMD connection contained in the pipeline_root directory with file name \"metadata.sqlite\" will be used.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PIPELINE_NAME = 'tfx-covertype-classifier'\n",
    "PIPELINE_ROOT = os.path.join(ARTIFACT_STORE, PIPELINE_NAME, time.strftime(\"%Y%m%d_%H%M%S\"))\n",
    "os.makedirs(PIPELINE_ROOT, exist_ok=True)\n",
    "\n",
    "context = InteractiveContext(\n",
    "    pipeline_name=PIPELINE_NAME,\n",
    "    pipeline_root=PIPELINE_ROOT,\n",
    "    metadata_connection_config=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ingesting data using ExampleGen\n",
    "\n",
    "In any ML development process the first step  is to ingest the training and test datasets. The `ExampleGen` component ingests data into a TFX pipeline. It consumes external files/services to generate a set file files in the `TFRecord` format,  which will be used by other TFX components. It can also shuffle the data and split into an arbitrary number of partitions.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/ExampleGen.png width=\"300\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run CsvExampleGen\n",
    "\n",
    "In this exercise, you use the `CsvExampleGen` specialization of `ExampleGen` to ingest CSV files from a GCS location and emit them as `tf.Example` records for consumption by downstream TFX pipeline components. Your task is to configure the component to create 80-20 `train` and `eval` splits. *Hint: review the [ExampleGen proto](https://github.com/tensorflow/tfx/blob/master/tfx/proto/example_gen.proto#L78) definition to split your data with hash buckets.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_config = example_gen_pb2.Output(\n",
    "    split_config=example_gen_pb2.SplitConfig(splits=[\n",
    "        # TODO: Your code to configure train data split\n",
    "        # TODO: Your code to configure eval data split        \n",
    "        example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=4),\n",
    "        example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1)\n",
    "    ]))\n",
    "\n",
    "example_gen = tfx.components.CsvExampleGen(\n",
    "    input_base=DATA_ROOT,\n",
    "    output_config=output_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(example_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the ingested data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples_uri = example_gen.outputs['examples'].get()[0].uri\n",
    "tfrecord_filenames = [os.path.join(examples_uri, 'train', name)\n",
    "                      for name in os.listdir(os.path.join(examples_uri, 'train'))]\n",
    "dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type=\"GZIP\")\n",
    "for tfrecord in dataset.take(2):\n",
    "  example = tf.train.Example()\n",
    "  example.ParseFromString(tfrecord.numpy())\n",
    "  for name, feature in example.features.feature.items():\n",
    "    if feature.HasField('bytes_list'):\n",
    "        value = feature.bytes_list.value\n",
    "    if feature.HasField('float_list'):\n",
    "        value = feature.float_list.value\n",
    "    if feature.HasField('int64_list'):\n",
    "        value = feature.int64_list.value\n",
    "    print('{}: {}'.format(name, value))\n",
    "  print('******')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating statistics using StatisticsGen\n",
    "\n",
    "The `StatisticsGen`  component generates data statistics that can be used by other TFX components. StatisticsGen uses [TensorFlow Data Validation](https://www.tensorflow.org/tfx/data_validation/get_started). `StatisticsGen` generate statistics for each split in the `ExampleGen` component's output. In our case there two splits: `train` and `eval`.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/StatisticsGen.png width=\"200\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and  run the `StatisticsGen` component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "statistics_gen = tfx.components.StatisticsGen(\n",
    "    examples=example_gen.outputs['examples'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(statistics_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize statistics\n",
    "\n",
    "The generated statistics can be visualized using the `tfdv.visualize_statistics()` function from the [TensorFlow Data Validation](https://www.tensorflow.org/tfx/data_validation/get_started) library or using a utility method of the `InteractiveContext` object. In fact, most of the artifacts generated by the TFX components can be visualized using `InteractiveContext`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(statistics_gen.outputs['statistics'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infering data schema using SchemaGen\n",
    "\n",
    "Some TFX components use a description input data called a schema. The schema is an instance of `schema.proto`. It can specify data types for feature values, whether a feature has to be present in all examples, allowed value ranges, and other properties. `SchemaGen` automatically generates the schema by inferring types, categories, and ranges from data statistics. The auto-generated schema is best-effort and only tries to infer basic properties of the data. It is expected that developers review and modify it as needed. `SchemaGen` uses [TensorFlow Data Validation](https://www.tensorflow.org/tfx/data_validation/get_started).\n",
    "\n",
    "The `SchemaGen` component generates the schema using the statistics for the `train` split. The statistics for other splits are ignored.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/SchemaGen.png width=\"200\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the `SchemaGen` components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_gen = SchemaGen(\n",
    "    statistics=statistics_gen.outputs['statistics'],\n",
    "    infer_feature_shape=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(schema_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the inferred schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(schema_gen.outputs['schema'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Updating the auto-generated schema\n",
    "\n",
    "In most cases the auto-generated schemas must be fine-tuned manually using insights from data exploration and/or domain knowledge about the data. For example, you know that in the `covertype` dataset there are seven types of forest cover (coded using 1-7 range) and that the value of the `Slope` feature should be in the 0-90 range. You can manually add these constraints to the auto-generated schema by setting the feature domain."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the auto-generated schema proto file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_proto_path = '{}/{}'.format(schema_gen.outputs['schema'].get()[0].uri, 'schema.pbtxt')\n",
    "schema = tfdv.load_schema_text(schema_proto_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modify the schema\n",
    "\n",
    "You can use the protocol buffer APIs to modify the schema. \n",
    "\n",
    "**Hint**: Review the [TFDV library API documentation](https://www.tensorflow.org/tfx/data_validation/api_docs/python/tfdv/set_domain) on setting a feature's domain. You can use the protocol buffer APIs to modify the schema. Review the [Tensorflow Metadata proto definition](https://github.com/tensorflow/metadata/blob/master/tensorflow_metadata/proto/v0/schema.proto#L405) for configuration options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Your code to restrict the categorical feature Cover_Type between the values of 0 and 6.\n",
    "# TODO: Your code to restrict the numeric feature Slope between 0 and 90.\n",
    "\n",
    "tfdv.set_domain(schema, 'Cover_Type', schema_pb2.IntDomain(name='Cover_Type', min=0, max=6, is_categorical=True))\n",
    "tfdv.set_domain(schema, 'Slope',  schema_pb2.IntDomain(name='Slope', min=0, max=90))\n",
    "\n",
    "tfdv.display_schema(schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the updated schema "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_dir = os.path.join(ARTIFACT_STORE, 'schema')\n",
    "tf.io.gfile.makedirs(schema_dir)\n",
    "schema_file = os.path.join(schema_dir, 'schema.pbtxt')\n",
    "\n",
    "tfdv.write_schema_text(schema, schema_file)\n",
    "\n",
    "!cat {schema_file}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the updated schema using ImporterNode\n",
    "\n",
    "The `ImporterNode` component allows you to import an external artifact, including the schema file, so it can be used by other TFX components in your workflow. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the `ImporterNode` component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_importer = ImporterNode(\n",
    "    instance_name='Schema_Importer',\n",
    "    source_uri=schema_dir,\n",
    "    artifact_type=tfx.types.standard_artifacts.Schema,\n",
    "    reimport=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(schema_importer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the imported schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(schema_importer.outputs['result'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validating data with ExampleValidator\n",
    "\n",
    "The `ExampleValidator` component identifies anomalies in data.  It identifies anomalies by comparing data statistics computed by the `StatisticsGen` component against a schema generated by `SchemaGen` or imported by `ImporterNode`.\n",
    "\n",
    "`ExampleValidator` can detect different classes of anomalies. For example it can:\n",
    "\n",
    "- perform validity checks by comparing data statistics against a schema \n",
    "- detect training-serving skew by comparing training and serving data.\n",
    "- detect data drift by looking at a series of data.\n",
    "\n",
    "\n",
    "The `ExampleValidator` component validates the data in the `eval` split only. Other splits are ignored. \n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/ExampleValidator.png width=\"350\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the `ExampleValidator` component\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Complete ExampleValidator\n",
    "# Hint: review the visual above and review the documentation on ExampleValidator's inputs and outputs: \n",
    "# https://www.tensorflow.org/tfx/guide/exampleval\n",
    "# Make sure you use the output of the schema_importer component created above.\n",
    "\n",
    "example_validator = ExampleValidator(\n",
    "    instance_name=\"Data_Validation\",    \n",
    "    statistics=statistics_gen.outputs['statistics'],\n",
    "    schema=schema_importer.outputs['result'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(example_validator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the output of `ExampleValidator`\n",
    "\n",
    "The output artifact of the ExampleValidator is the `anomalies.pbtxt` file describing an anomalies_pb2.Anomalies protobuf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_uri = example_validator.outputs['anomalies'].get()[0].uri\n",
    "train_anomalies_filename = os.path.join(train_uri, \"train/anomalies.pbtxt\")\n",
    "!cat $train_anomalies_filename"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize validation results\n",
    "\n",
    "The file `anomalies.pbtxt` can be visualized using `context.show`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(example_validator.outputs['output'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our case no anomalies were detected in the `eval` split.\n",
    "\n",
    "For a detailed deep dive into data validation and schema generation refer to the `lab-31-tfdv-structured-data` lab."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing data with Transform\n",
    "\n",
    "The `Transform` component performs data transformation and feature engineering. The `Transform` component consumes `tf.Examples` emitted from the `ExampleGen` component and emits the transformed feature data and the `SavedModel` graph that was used to process the data. The emitted `SavedModel`  can then be used by serving components to make sure that the same data pre-processing logic is applied at training and serving.\n",
    "\n",
    "The `Transform` component requires more code than many other components because of the arbitrary complexity of the feature engineering that you may need for the data and/or model that you're working with. It requires code files to be available which define the processing needed.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Transform.png width=\"400\">\n",
    "\n",
    "### Define the pre-processing module\n",
    "\n",
    "To configure `Trainsform`, you need to encapsulate your pre-processing code in the Python `preprocessing_fn` function and save it to a  python module that is then provided to the Transform component as an input. This module will be loaded by transform and the `preprocessing_fn` function will be called when the `Transform` component runs.\n",
    "\n",
    "In most cases, your implementation of the `preprocessing_fn` makes extensive use of [TensorFlow Transform](https://www.tensorflow.org/tfx/guide/tft) for performing feature engineering on your dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRANSFORM_MODULE = 'preprocessing.py'\n",
    "!cat {TRANSFORM_MODULE}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the `Transform` component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = Transform(\n",
    "    examples=example_gen.outputs['examples'],\n",
    "    schema=schema_importer.outputs['result'],\n",
    "    module_file=TRANSFORM_MODULE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the `Transform` component's outputs\n",
    "\n",
    "The Transform component has 2 outputs:\n",
    "\n",
    "- `transform_graph` - contains the graph that can perform the preprocessing operations (this graph will be included in the serving and evaluation models).\n",
    "- `transformed_examples` - contains the preprocessed training and evaluation data.\n",
    "\n",
    "Take a peek at the `transform_graph` artifact: it points to a directory containing 3 subdirectories:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir(transform.outputs['transform_graph'].get()[0].uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the `transform.examples` artifact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir(transform.outputs['transformed_examples'].get()[0].uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_uri = transform.outputs['transformed_examples'].get()[0].uri\n",
    "tfrecord_filenames = [os.path.join(transform_uri,  'train', name)\n",
    "                      for name in os.listdir(os.path.join(transform_uri, 'train'))]\n",
    "dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type=\"GZIP\")\n",
    "for tfrecord in dataset.take(2):\n",
    "  example = tf.train.Example()\n",
    "  example.ParseFromString(tfrecord.numpy())\n",
    "  for name, feature in example.features.feature.items():\n",
    "    if feature.HasField('bytes_list'):\n",
    "        value = feature.bytes_list.value\n",
    "    if feature.HasField('float_list'):\n",
    "        value = feature.float_list.value\n",
    "    if feature.HasField('int64_list'):\n",
    "        value = feature.int64_list.value\n",
    "    print('{}: {}'.format(name, value))\n",
    "  print('******')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train your TensorFlow model with the `Trainer` component\n",
    "\n",
    "The `Trainer` component trains a model using TensorFlow.\n",
    "\n",
    "`Trainer` takes:\n",
    "\n",
    "- tf.Examples used for training and eval.\n",
    "- A user provided module file that defines the trainer logic.\n",
    "- A data schema created by `SchemaGen` or imported by `ImporterNode`.\n",
    "- A proto definition of train args and eval args.\n",
    "- An optional transform graph produced by upstream Transform component.\n",
    "- An optional base models used for scenarios such as warmstarting training.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Trainer.png width=\"400\">\n",
    "\n",
    "\n",
    "### Define the trainer module\n",
    "\n",
    "To configure `Trainer`, you need to encapsulate your training code in a Python module that is then provided to the `Trainer` as an input. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINER_MODULE_FILE = 'model.py'\n",
    "!cat {TRAINER_MODULE_FILE}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and run the Trainer component\n",
    "\n",
    "As of the 0.25.0 release of TFX, the `Trainer` component only supports passing a single field - `num_steps` - through the `train_args` and `eval_args` arguments. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    custom_executor_spec=executor_spec.ExecutorClassSpec(trainer_executor.GenericExecutor),\n",
    "    module_file=TRAINER_MODULE_FILE,\n",
    "    transformed_examples=transform.outputs.transformed_examples,\n",
    "    schema=schema_importer.outputs.result,\n",
    "    transform_graph=transform.outputs.transform_graph,\n",
    "    train_args=trainer_pb2.TrainArgs(splits=['train'], num_steps=5000),\n",
    "    eval_args=trainer_pb2.EvalArgs(splits=['eval'], num_steps=1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyzing training runs with TensorBoard\n",
    "\n",
    "In this step you will analyze the training run with [TensorBoard.dev](https://blog.tensorflow.org/2019/12/introducing-tensorboarddev-new-way-to.html). `TensorBoard.dev` is a managed service that enables you to easily host, track and share your ML experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Retrieve the location of TensorBoard logs\n",
    "\n",
    "Each model run's train and eval metric logs are written to the `model_run` directory by the Tensorboard callback defined in `model.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logs_path = trainer.outputs['model_run'].get()[0].uri\n",
    "print(logs_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Upload the logs and start TensorBoard.dev\n",
    "\n",
    "1. Open a new JupyterLab terminal window\n",
    "\n",
    "2. From the terminal window, execute the following command\n",
    "```\n",
    "tensorboard dev upload --logdir [YOUR_LOGDIR]\n",
    "```\n",
    "\n",
    "Where [YOUR_LOGDIR] is an URI retrieved by the previous cell.\n",
    "\n",
    "You will be asked to authorize `TensorBoard.dev` using your Google account. If you don't have a Google account or you don't want to authorize `TensorBoard.dev` you can skip this exercise.\n",
    "\n",
    "After the authorization process completes, follow the link provided to view your experiment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tune your model's hyperparameters with the `Tuner` component"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Tuner` component makes use of the Python [KerasTuner](https://keras-team.github.io/keras-tuner/) API to tune your model's hyperparameters. It tighty integrates with the `Transform` and `Trainer` components for model hyperparameter tuning in continuous training pipelines as well as advanced use cases such as feature selection, feature engineering, and model architecture search.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Tuner_Overview.png width=\"400\">\n",
    "\n",
    "`Tuner` takes:\n",
    "\n",
    "*  A user provided module file (or module fn) that defines the tuning logic, including model definition, hyperparameter search space, objective etc.\n",
    "*  tf.Examples used for training and eval.\n",
    "*  Protobuf definition of train args and eval args.\n",
    "  * (Optional) Protobuf definition of tuning args.\n",
    "  * (Optional) transform graph produced by an upstream Transform component.\n",
    "  * (Optional) A data schema created by a SchemaGen pipeline component and optionally altered by the developer.\n",
    "  \n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Tuner.png width=\"400\">\n",
    "\n",
    "With the given data, model, and objective, `Tuner` tunes the hyperparameters and emits the best results that can be directly fed into the `Trainer` component during model re-training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner = Tuner(\n",
    "        module_file=TRAINER_MODULE_FILE,\n",
    "        examples=transform.outputs['transformed_examples'],\n",
    "        transform_graph=transform.outputs['transform_graph'],\n",
    "        train_args=trainer_pb2.TrainArgs(num_steps=1000),\n",
    "        eval_args=trainer_pb2.EvalArgs(num_steps=500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(tuner)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrain your model by running Tuner with the best hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams_importer = ImporterNode(\n",
    "    instance_name='import_hparams',\n",
    "    # This can be Tuner's output file or manually edited file. The file contains\n",
    "    # text format of hyperparameters (kerastuner.HyperParameters.get_config())\n",
    "    source_uri=tuner.outputs.best_hyperparameters.get()[0].uri,\n",
    "    artifact_type=HyperParameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(hparams_importer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: your code to retrain your model with the best hyperparameters found by the Tuner component above.\n",
    "# Hint: review the Trainer code above in this notebook and the documentation for how to configure the trainer to use the output artifact from the hparams_importer.\n",
    "\n",
    "trainer = Trainer(\n",
    "    custom_executor_spec=executor_spec.ExecutorClassSpec(trainer_executor.GenericExecutor),\n",
    "    module_file=TRAINER_MODULE_FILE,\n",
    "    transformed_examples=transform.outputs.transformed_examples,\n",
    "    schema=schema_importer.outputs.result,\n",
    "    transform_graph=transform.outputs.transform_graph,\n",
    "    hyperparameters=hparams_importer.outputs.result,    \n",
    "    train_args=trainer_pb2.TrainArgs(splits=['train'], num_steps=5000),\n",
    "    eval_args=trainer_pb2.EvalArgs(splits=['eval'], num_steps=1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating trained models with Evaluator\n",
    "The `Evaluator` component analyzes model performance using the [TensorFlow Model Analysis library](https://www.tensorflow.org/tfx/model_analysis/get_started). It runs inference requests on particular subsets of the test dataset, based on which slices are defined by the developer. Knowing which slices should be analyzed requires domain knowledge of what is important in this particular use case or domain. \n",
    "\n",
    "The `Evaluator` can also optionally validate a newly trained model against a previous model. In this lab, you only train one model, so the Evaluator automatically will label the model as \"blessed\".\n",
    "\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Evaluator.png width=\"400\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the Evaluator component\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the `ResolverNode` to pick the previous model to compare against.  The model resolver is only required if performing model validation in addition to evaluation. In this case we validate against the latest blessed model. If no model has been blessed before (as in this case) the evaluator will make our candidate the first blessed model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resolver = ResolverNode(\n",
    "      instance_name='latest_blessed_model_resolver',\n",
    "      resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,\n",
    "      model=Channel(type=Model),\n",
    "      model_blessing=Channel(type=ModelBlessing))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(model_resolver)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Configure evaluation metrics and slices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Your code here to create a tfma.MetricThreshold. \n",
    "# Review the API documentation here: https://www.tensorflow.org/tfx/model_analysis/api_docs/python/tfma/MetricThreshold\n",
    "# Hint: Review the API documentation for tfma.GenericValueThreshold to constrain accuracy between 50% and 99%.\n",
    "\n",
    "accuracy_threshold = tfma.MetricThreshold(\n",
    "                value_threshold=tfma.GenericValueThreshold(\n",
    "                    lower_bound={'value': 0.5},\n",
    "                    upper_bound={'value': 0.99})\n",
    "                )\n",
    "\n",
    "metrics_specs = tfma.MetricsSpec(\n",
    "                   metrics = [\n",
    "                       tfma.MetricConfig(class_name='SparseCategoricalAccuracy',\n",
    "                           threshold=accuracy_threshold),\n",
    "                       tfma.MetricConfig(class_name='ExampleCount')])\n",
    "\n",
    "eval_config = tfma.EvalConfig(\n",
    "    model_specs=[\n",
    "        tfma.ModelSpec(label_key='Cover_Type')\n",
    "    ],\n",
    "    metrics_specs=[metrics_specs],\n",
    "    slicing_specs=[\n",
    "        tfma.SlicingSpec(),\n",
    "        tfma.SlicingSpec(feature_keys=['Wilderness_Area'])\n",
    "    ]\n",
    ")\n",
    "eval_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_analyzer = Evaluator(\n",
    "    examples=example_gen.outputs.examples,\n",
    "    model=trainer.outputs.model,\n",
    "    baseline_model=model_resolver.outputs.model,\n",
    "    eval_config=eval_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(model_analyzer, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the model performance validation status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_blessing_uri = model_analyzer.outputs.blessing.get()[0].uri\n",
    "!ls -l {model_blessing_uri}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize evaluation results\n",
    "You can visualize the evaluation results using the `tfma.view.render_slicing_metrics()` function from TensorFlow Model Analysis library.\n",
    "\n",
    "**Setup Note:** *Currently, TFMA visualizations don't render in  JupyterLab. Make sure that you run this notebook in Classic Notebook.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_uri = model_analyzer.outputs['evaluation'].get()[0].uri\n",
    "evaluation_uri\n",
    "!ls {evaluation_uri}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_result = tfma.load_eval_result(evaluation_uri)\n",
    "eval_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfma.view.render_slicing_metrics(eval_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfma.view.render_slicing_metrics(\n",
    "    eval_result, slicing_column='Wilderness_Area')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## InfraValidator\n",
    "\n",
    "The `InfraValidator` component acts as an additional early warning layer by validating a candidate model in a sandbox version of its serving infrastructure to prevent an unservable model from being pushed to production. Compared to the `Evaluator` component above which validates a model's performance, the `InfraValidator` component is validating that a model is able to generate predictions from served examples in an environment configured to match production. The config below takes a model and examples, launches the model in a sand-boxed [TensorflowServing](https://www.tensorflow.org/tfx/guide/serving) model server from the latest image in a local docker engine, and optionally checks that the model binary can be loaded and queried before \"blessing\" it for production.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/InfraValidator.png width=\"400\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infra_validator = InfraValidator(\n",
    "    model=trainer.outputs['model'],\n",
    "    examples=example_gen.outputs['examples'],\n",
    "    serving_spec=infra_validator_pb2.ServingSpec(\n",
    "        tensorflow_serving=infra_validator_pb2.TensorFlowServing(\n",
    "            tags=['latest']),\n",
    "      local_docker=infra_validator_pb2.LocalDockerConfig(),\n",
    "  ),\n",
    "    validation_spec=infra_validator_pb2.ValidationSpec(\n",
    "        max_loading_time_seconds=60,\n",
    "        num_tries=5,\n",
    "    ),    \n",
    "  request_spec=infra_validator_pb2.RequestSpec(\n",
    "      tensorflow_serving=infra_validator_pb2.TensorFlowServingRequestSpec(),\n",
    "          num_examples=5,\n",
    "      )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.run(infra_validator, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the model infrastructure validation status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infra_blessing_uri = infra_validator.outputs.blessing.get()[0].uri\n",
    "!ls -l {infra_blessing_uri}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploying models with Pusher\n",
    "\n",
    "The `Pusher` component checks whether a model has been \"blessed\", and if so, deploys it by pushing the model to a well known file destination.\n",
    "\n",
    "<img src=https://github.com/GoogleCloudPlatform/mlops-on-gcp/raw/master/images/Pusher.png width=\"400\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and run the `Pusher` component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.outputs['model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pusher = Pusher(\n",
    "    model=trainer.outputs['model'],\n",
    "    model_blessing=model_analyzer.outputs['blessing'],\n",
    "    infra_blessing=infra_validator.outputs['blessing'],\n",
    "    push_destination=pusher_pb2.PushDestination(\n",
    "        filesystem=pusher_pb2.PushDestination.Filesystem(\n",
    "            base_directory=SERVING_MODEL_DIR)))\n",
    "context.run(pusher)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine the output of `Pusher`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pusher.outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set `PATH` to include a directory containing `saved_model_cli.\n",
    "PATH=%env PATH\n",
    "%env PATH=/opt/conda/envs/tfx/bin:{PATH}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latest_pushed_model = os.path.join(SERVING_MODEL_DIR, max(os.listdir(SERVING_MODEL_DIR)))\n",
    "!saved_model_cli show --dir {latest_pushed_model} --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "This concludes your introductory walthrough through TFX pipeline components. In the lab, you used TFX to analyze, understand, and pre-process the dataset and train, analyze, validate, and deploy a multi-class classification model to predict the type of forest cover from cartographic features. You utilized a TFX Interactive Context for prototype development of a TFX pipeline directly in a Jupyter notebook. Next, you worked with the TFDV library to modify your dataset schema to add feature constraints to catch data anamolies that can negatively impact your model's performance. You utilized TFT library for feature proprocessing for consistent feature transformations for your model at training and serving time. Lastly, using the TFMA library, you added model performance constraints to ensure you only push more accurate models than previous runs to production.\n",
    "\n",
    "\n",
    "The next labs in the series will guide through developing a TFX pipeline, deploying and running the pipeline on **AI Platform Pipelines** and automating the pipeline build and deployment processes with **Cloud Build**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## License"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font size=-1>Licensed under the Apache License, Version 2.0 (the \\\"License\\\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \\\"AS IS\\\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the License for the specific language governing permissions and limitations under the License.</font>"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-2-3-gpu.2-3.m59",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-2-3-gpu.2-3:m59"
  },
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
